# Copyright 2023 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse
import os

from acl_resource import AclResource
from acl_model import Model
from kws_utils import txt_to_list, infer_wav
    

TRUST_VAL = 0.8

    
def parse_args():
    parser = argparse.ArgumentParser(description='Your script description')
    parser.add_argument("--mode", choices=["single", "dataset"], required=True, help="Specify the mode ('single' or 'dataset').")
    parser.add_argument('--model_path', type=str, help='Path to the model file (om)', default='./models/kws.om') 
    parser.add_argument('--test_wav', type=str, help='(mode single) Testing .wav file', default='./keyword.wav') 
    parser.add_argument('--folder_path', type=str, help='(mode dataset) Testing .wav folder', default="./dataset_shuffle") 
    parser.add_argument('--gt_file', type=str, help='(mode dataset) path of ground truth file', default="./dataset_shuffle/label.txt")
    args = parser.parse_args()
    return args


def check_args(args):
    if not os.path.exists(args.model_path):
        raise FileNotFoundError(f"Model file {args.model_path} not found.")
    if args.mode == "single" and not os.path.exists(args.test_wav):
        raise FileNotFoundError(f"Test wav file {args.test_wav} not found.")
    if args.mode == "dataset" and not os.path.exists(args.folder_path):
        raise FileNotFoundError(f"Test folder {args.folder_path} not found.")
    if args.mode == "dataset" and not os.path.exists(args.gt_file):
        raise FileNotFoundError(f"Ground truth file {args.gt_file} not found.")


def single_infer(model, test_wav):
    keyword_score = infer_wav(test_wav, model)
    print(f"Keyword score is: {keyword_score}")
    if keyword_score > TRUST_VAL:
        print("this wav is waked up: True")
    else:
        print("this wav is waked up: False")
            
    
def dataset_infer(model, test_folder, gt_file):
    correct_num = 0
    gt_list = txt_to_list(gt_file)    
    
    gt_idx = 0
    list_files = os.listdir(test_folder)
    list_files.remove('label.txt')
    list_files.sort(key=lambda x: int(x[:-4]))
    for filename in list_files:
        if filename.endswith(".wav"):
            wav_path = os.path.join(test_folder, filename)
            keyword_score = infer_wav(wav_path, model)
            keyword_det = 1 if keyword_score > TRUST_VAL else 0
            if keyword_det == gt_list[gt_idx]:
                correct_num += 1
            gt_idx += 1
            print(f"Keyword score for {filename}: {keyword_det}")
    print("correct_num: ", correct_num)
    print("total_num: ", len(gt_list))
    print("accuracy: ", correct_num / len(gt_list))

    
def main(model, test_wav, mode, test_folder, gt_file):
    if mode == "single":
        single_infer(model, test_wav)
    else:
        dataset_infer(model, test_folder, gt_file)


if __name__ == "__main__":
    args = parse_args()
    check_args(args)
    model_path = args.model_path
    test_wav = args.test_wav
    mode = args.mode
    test_folder = args.folder_path
    gt_file = args.gt_file
    
    acl_resource = AclResource()
    acl_resource.init()
    model = Model(acl_resource, model_path)
    main(model, test_wav, mode, test_folder, gt_file)
